from PyPDF2 import PdfReader, PdfWriter
from datetime import datetime
import os

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

def safe_extract_pages(page_array=None, txt_path=None,input_path="origin.pdf", output_path=f"data_{timestamp}.pdf"):
    try:
        # 验证文件存在
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件 {input_path} 不存在")

        if page_array is None and txt_path is None:
            raise ValueError("必须提供 page_array 或 txt_path 参数")

        # 优先级：直接输入 > 文件读取
        if page_array is not None:
            pages = [int(p)-1 for p in page_array]  # 1-based转0-based
        else:
            if not os.path.exists(txt_path):
                raise FileNotFoundError(f"输入页码文件 {txt_path} 不存在")
            with open(txt_path, 'r', encoding='utf-8') as f:
                pages = [int(line.strip())-1 for line in f if line.strip().isdigit()]

        reader = PdfReader(input_path)
        writer = PdfWriter()
        total_pages = len(reader.pages)

        # 提取页面并验证页码
        extracted_pages = 0
        for page_num in pages:
            if 0 <= page_num < total_pages:
                writer.add_page(reader.pages[page_num])
                extracted_pages += 1
            else:
                print(f"警告：跳过无效页码 {page_num+1}（总页数：{total_pages}）")

        # 保存结果
        with open(output_path, "wb") as f:
            writer.write(f)
        print(f"成功提取 {extracted_pages} 页到 {os.path.abspath(output_path)}")

    except Exception as e:
        print(f"处理失败：{str(e)}")


def process_mode1():
    pdf_path_str = input("请输入要处理的PDF路径及名称:\n")
    input_str = input("请输入需要提取的页码（用逗号或空格分隔:\n")
    txt_path_str = input("请输入要处理的PDF页码路径及名称:\n")
    page_array = None
    if (input_str != ""):
        page_array = [int(x) for x in input_str.replace(',', ' ').split()]
    filename = "Processed\\" + input("请为输出文件命名(仅支持下划线)") + ".pdf"
    if (txt_path_str == ""):
        txt_path_str = "page_array.txt"

    safe_extract_pages(page_array, txt_path_str, pdf_path_str, )

def process_mode2():
    pdf_path_str = input("请输入要处理的PDF路径及名称:\n")
    input_str = input("请输入需要提取的页码（用逗号或空格分隔:\n")
    txt_path_str = input("请输入要处理的PDF页码路径及名称:\n")
    page_array = None
    if (input_str != ""):
        page_array = [int(x) for x in input_str.replace(',', ' ').split()]
    filename = "Processed\\" + input("请为输出文件命名(仅支持下划线)") + ".pdf"
    if (txt_path_str == ""):
        txt_path_str = "page_array.txt"

    safe_extract_pages(page_array, txt_path_str, pdf_path_str, txt_path_str)



# 使用示例（页码从0开始计数）
if __name__ == "__main__":
    # 示例：提取第1、3、5页（对应页码0, 2, 4）
    process_mode1()
